# -*- coding: utf-8 -*-

# Copyright (c) JoinQuant Development Team
# Author: Huayong Kuang <kuanghuayong@joinquant.com>

import zmq
from multiprocessing import Process

def calculate(data):
    return {"output": data['ok']}

def handle_request(url):
    context = zmq.Context()
    socket = context.socket(zmq.REP)
    socket.connect(url)
    while True:
        data = socket.recv_json()
        print("received {}".format(data))
        out = calculate(data)
        socket.send_json(out)

def server():
    context = zmq.Context()
    # Set up socket for clients to connect to.
    clients = context.socket(zmq.ROUTER)
    clients.bind('tcp://*:{}'.format(5556))

    # Set up ipc socket for workers to connect to
    url_worker = 'ipc:///tmp/workers'
    workers = context.socket(zmq.DEALER)
    workers.bind(url_worker)

    # Start 4 worker processes
    for _ in range(4):
        p = Process(target=handle_request, args=(url_worker,))
        p.start()

    # Forward requests from clients to the workers via a Queue
    zmq.device(zmq.QUEUE, clients, workers)


if __name__ == "__main__":
    server()

